package cn.edu.thu.tsquality.core.algorithm.auto.regression;

import cn.edu.thu.tsquality.core.common.algorithm.IAlgorithm;
import cn.edu.thu.tsquality.core.common.table.Header;
import cn.edu.thu.tsquality.core.common.table.Row;
import cn.edu.thu.tsquality.core.common.table.Table;
import javafx.util.Pair;
import weka.classifiers.functions.LinearRegression;
import weka.core.Attribute;
import weka.core.DenseInstance;
import weka.core.Instance;
import weka.core.Instances;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/*
自回归修复不完整/不精确的值
修复算法如下：
若needRepair[i] == true
那么使用attrData[i-regressionSize], attrData[i-(regressionSize-1)], ..., attrData[i-1]回归得到attrData[i]
*/

public class AutoRegression implements IAlgorithm
{
    private Integer regressionSize;

    AutoRegression(Integer regressionSize)
    {
        this.regressionSize = regressionSize;
    }


    private void println(String str)
    {
        System.out.println(str);
    }

    private void print(String str)
    {
        System.out.print(str);
    }

    public Table run(Table table)
    {
        List tableAttrList = Arrays.asList(table.getHeader().toArray());
        if (tableAttrList.size() !=1 && tableAttrList.size()!= 2)
            throw new IllegalArgumentException("table columns error, expected 2 columns");

        List<Row> tableRowsList = table.getRows();
        int inputDataSize = tableRowsList.size();

        double[] attrData = new double[inputDataSize];
        boolean[] needRepair = new boolean[inputDataSize];

        // 求需要修复的点
        for (int i = 0; i < inputDataSize; i++)
        {
            if(tableRowsList.get(i).get(0).equals(""))
            {
                attrData[i] = 0.0;
                needRepair[i] = true;
            }
            else
            {
                attrData[i] = Double.parseDouble((String) tableRowsList.get(i).get(0));

                if(table.getHeader().size() == 2 && Integer.parseInt((String) tableRowsList.get(i).get(1)) != 0)
                    needRepair[i] = true;
                else
                    needRepair[i] = false;
            }
        }


        // 返回值
        Header returnHeader = new Header(new String[] { (String)tableAttrList.get(0), "cleaned_"+tableAttrList.get(0)});
        List<Row> returnRows = new ArrayList<>();

        // 自回归训练集，由ArrayList<Double> 回归 Double
        ArrayList<Pair<ArrayList<Double>, Double>> dataSet = new ArrayList<>();
        for(int i = 0; i < inputDataSize; i++)
        {
            boolean isInstance = true;
            for(int j = 0; j <= regressionSize; j++)
            {
                if(i + j >= inputDataSize || needRepair[i + j])
                {
                    isInstance = false;
                    break;
                }
            }
            if(isInstance)
            {
                ArrayList<Double> X = new ArrayList<>();
                for(int j = 0; j < regressionSize; j++)
                    X.add(attrData[i+j]);
                dataSet.add(new Pair<>(X, attrData[i+regressionSize]));
            }
        }

        if(dataSet.size() > 0)
        {
            LinearRegression linearRegression = new LinearRegression();
            ArrayList<Attribute> attrs = new ArrayList<>();
            for (int i = 0; i <= regressionSize; i++)
                attrs.add(new Attribute(""+i));
            Instances instances = new Instances("dataSet", attrs, dataSet.size());
            instances.setClassIndex(regressionSize);

            for(Pair<ArrayList<Double>, Double> pair : dataSet)
            {
                Instance instance = new DenseInstance(regressionSize + 1);
                instance.setDataset(instances);
                for(int i = 0; i < pair.getKey().size(); i++)
                    instance.setValue(attrs.get(i), pair.getKey().get(i));
                instance.setValue(attrs.get(regressionSize), pair.getValue());
                instances.add(instance);
            }

            try
            {
                linearRegression.buildClassifier(instances);
            }
            catch (Exception e)
            {
                System.out.println("train error! "+e.toString());
                return new Table(returnHeader, returnRows);
            }

            //------------
//            println(regressionSize+"");
//            println(inputDataSize+"");
//            print("data = [");
//            for(Double i : attrData)
//                print(i+" ");
//            println("]");

//            int coefIndex = 0;
//            println("y["+regressionSize+"]=");
//            for(double coef : linearRegression.coefficients())
//                println(coef + (coefIndex==regressionSize+1?"": " * y[" + coefIndex++ +"] + "));

            // 进行修复
            for (int i = 0; i < inputDataSize; i++)
            {
                if (needRepair[i])
                {
                    Instance instance = new DenseInstance(regressionSize);
                    for(int j = 0; j < regressionSize; j++)
                    {
                        double data = i - regressionSize + j >= 0? attrData[i-regressionSize+j] : attrData[0];
                        instance.setValue(attrs.get(j), data);
                    }

                    try
                    {
                        double repairedValue = linearRegression.classifyInstance(instance);
                        attrData[i] = repairedValue;
                    }
                    catch (Exception e)
                    {
                        System.out.println("test error! "+e.toString());
                        attrData[i] = 0.0;
                    }
                }

                returnRows.add(new Row(returnHeader, new String[]{
                        (String) tableRowsList.get(i).get(0),
                        String.format("%.3f", attrData[i])}, i));
            }
        }
        return new Table(returnHeader, returnRows);
    }
}
